package org.apache.solr.ocean.search;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;

import org.apache.lucene.search.FieldDocSortedHitQueue;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.Sort;
import org.apache.solr.ocean.core.Snapshot;
import org.apache.solr.ocean.core.Index.IndexSnapshot;
import org.apache.solr.ocean.util.Timeout;

public class SnapshotSearch extends Search {
  private Snapshot snapshot;
  private ExecutorService executorService;
  private long startTime;

  public SnapshotSearch(SearchQuery searchQuery, Snapshot snapshot, ExecutorService executorService) {
    super(searchQuery);
    this.snapshot = snapshot;
    this.executorService = executorService;
  }

  public Result call() throws Exception {
    startTime = Timeout.TIMER_THREAD.getTime();
    Collection<IndexSnapshot> indexSnapshots = snapshot.getIndexSnapshots();
    List<Callable<Result>> indexSearches = new ArrayList<Callable<Result>>(indexSnapshots.size());
    for (IndexSnapshot indexSnapshot : indexSnapshots) {
      IndexSnapshotSearch indexSearch = new IndexSnapshotSearch(searchQuery, indexSnapshot);
      indexSearches.add(indexSearch);
    }
    List<Future<Result>> futures = executorService.invokeAll(indexSearches);
    List<Result> results = new ArrayList<Result>(indexSnapshots.size());
    for (Future<Result> future : futures) {
      Result result = future.get();
      results.add(result);
    }
    if (searchQuery.hasSort()) {
      return mergeSort(searchQuery.getSort(), results);
    } else {
      return mergeSort(Sort.RELEVANCE, results);
    }
  }

  private Result mergeSort(Sort sort, List<Result> results) throws Exception {
    int numDocs = searchQuery.getNumDocs();
    float maxScore = Float.NEGATIVE_INFINITY;
    FieldDocSortedHitQueue fieldSortedQueue = new FieldDocSortedHitQueue(sort.getSort(), numDocs);
    int numFound = 0;
    for (Result result : results) {
      numFound += result.getNumFound();
      ScoreDoc[] scoreDocs = result.getAllDocs();
      maxScore = Math.max(maxScore, result.getMaxScore());
      for (ScoreDoc scoreDoc : scoreDocs) {
        fieldSortedQueue.insert(scoreDoc);
      }
    }
    ScoreDoc[] mergedScoreDocs = toArray(fieldSortedQueue);
    long elapsed = Timeout.TIMER_THREAD.getTime() - startTime;
    Result mergedResult = new Result(searchQuery.getStart(), mergedScoreDocs, numFound, elapsed, maxScore);
    return mergedResult;
  }
}
